import numpy as np
import torch
import torchvision.transforms as transforms
import yaml
from PIL import Image

from ..configs import read_config_file
from . import build_model

TRANSFORMER = transforms.Compose(
    [
        transforms.Resize((480, 640)),  # (height, width)
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.48501960784313836, 0.4579568627450961, 0.4076039215686255],
            std=[0.00392156862745098, 0.00392156862745098, 0.00392156862745098],
        ),
    ]
)


def extract_netvlad(img_path: str):
    model = build_model.get_model()
    img = Image.open(img_path).convert("RGB")
    img = TRANSFORMER(img)
    img = img.cuda()

    with torch.no_grad():
        des = model(img.unsqueeze(0))[0]

    des = des.cpu().numpy()

    return des
